#include <jni.h>
#include <string>
#include <unistd.h>

#include "hashmap.h"

extern "C" {
#include "libswresample/swresample.h"
#include "libavutil/opt.h"
#include "libavformat/avformat.h"
#include "libavcodec/mediacodec.h"
}

char *jstringToChar(JNIEnv *env, jstring string) {
    jsize charLen = env->GetStringLength(string);
    jsize utfLen = env->GetStringUTFLength(string);
    char *c = new char[utfLen + 1];
    env->GetStringUTFRegion(string, 0, charLen, c);
    return c;
}

extern "C"
JNIEXPORT jobject JNICALL
Java_com_zxhhyj_ffmpeg_tag_MediaTag_getMediaFormatStreamInfo(JNIEnv *env, jobject, jstring path) {
    const char *file_path = jstringToChar(env, path);
    AVFormatContext *formatCtx = avformat_alloc_context();
    if (avformat_open_input(&formatCtx, file_path, NULL, NULL) != 0) {
        delete[] file_path;
        avformat_close_input(&formatCtx);
        return nullptr;
    }
    if (avformat_find_stream_info(formatCtx, NULL) < 0) {
        delete[] file_path;
        avformat_close_input(&formatCtx);
        return nullptr;
    }
    av_dump_format(formatCtx, 0, file_path, 0);
    delete[] file_path;
    jlong duration = formatCtx->duration;
    jlong bitrate = formatCtx->bit_rate;
    int bitsPerRawSample = 0;
    jlong sampleRate = 0;
    int channels = 0;
    hashmap metadataMap = hashmap(env);
    AVDictionaryEntry *entry = nullptr;
    while ((entry = av_dict_get(formatCtx->metadata, "", entry, AV_DICT_IGNORE_SUFFIX))) {
        const char *key = entry->key;
        const char *value = entry->value;
        metadataMap.put(env->NewStringUTF(key), env->NewStringUTF(value));
    }
    for (int i = 0; i < formatCtx->nb_streams; i++) {
        if (formatCtx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {
            AVCodecParameters *codecParameters = formatCtx->streams[i]->codecpar;
            bitsPerRawSample = codecParameters->bits_per_raw_sample;
            sampleRate = codecParameters->sample_rate;
            channels = codecParameters->ch_layout.nb_channels;
            break;
        }
    }
    jbyteArray coverByteArray = NULL;
    for (int i = 0; i < formatCtx->nb_streams; ++i) {
        if (formatCtx->streams[i]->disposition & AV_DISPOSITION_ATTACHED_PIC) {
            std::vector<uint8_t> coverVector;
            const AVPacket packet = formatCtx->streams[i]->attached_pic;
            coverVector.insert(coverVector.end(), packet.data, packet.data + packet.size);
            coverByteArray = env->NewByteArray(coverVector.size());
            env->SetByteArrayRegion(coverByteArray, 0, coverVector.size(),
                                    reinterpret_cast<jbyte *>(coverVector.data()));
        }
    }
    avformat_close_input(&formatCtx);
    jclass formatSteamInfoClass = env->FindClass("com/zxhhyj/ffmpeg/tag/FormatStreamInfo");
    jmethodID formatSteamConstructor = env->GetMethodID(formatSteamInfoClass, "<init>",
                                                        "([BJJIJILjava/util/Map;)V");
    return env->NewObject(formatSteamInfoClass, formatSteamConstructor,
                          coverByteArray,
                          duration,
                          bitrate,
                          bitsPerRawSample,
                          sampleRate,
                          channels,
                          metadataMap.getJavaObject());
}

class PCMDecoder {
private:
    const std::vector<char> empty;
    char *path;
public:
    PCMDecoder();

    ~PCMDecoder();

    void init(const char *orgFilePath);

    std::vector<char> decode();

    void release();
};

PCMDecoder::PCMDecoder() {

}

PCMDecoder::~PCMDecoder() {

}

void PCMDecoder::init(const char *orgFilePath) {
    path = strdup(orgFilePath);
}

std::vector<char> PCMDecoder::decode() {
    AVFormatContext *formatCtx = avformat_alloc_context();
    AVPacket packet;
    AVFrame *pFrame = av_frame_alloc();
    int dataSize;
    jbyte *buffer;
    if (avformat_open_input(&formatCtx, path, NULL, NULL) != 0) {
        avformat_close_input(&formatCtx);
        return empty;
    }
    if (avformat_find_stream_info(formatCtx, NULL) < 0) {
        avformat_close_input(&formatCtx);
        return empty;
    }
    av_dump_format(formatCtx, 0, path, 0);
    const AVCodec *codec;
    AVCodecContext *codecContext = nullptr;
    int audioStreamIdx = av_find_best_stream(formatCtx, AVMEDIA_TYPE_AUDIO, -1, -1, NULL, 0);
    if (audioStreamIdx >= 0) {
        AVCodecParameters *codecParameters = formatCtx->streams[audioStreamIdx]->codecpar;
        codec = avcodec_find_decoder(codecParameters->codec_id);
        codecContext = avcodec_alloc_context3(codec);
    } else {
        return empty;
    }
    if (avcodec_open2(codecContext, codec, NULL) < 0) {
        avcodec_free_context(&codecContext);
        avformat_close_input(&formatCtx);
        return empty;
    }
// 分配 PCM 数据缓冲区
    dataSize = av_samples_get_buffer_size(NULL, codecContext->ch_layout.nb_channels,
                                          codecContext->frame_size, codecContext->sample_fmt, 1);
    buffer = (jbyte *) av_malloc(dataSize * sizeof(jbyte));

    // 循环解码音频数据
    while (av_read_frame(formatCtx, &packet) >= 0) {
        if (packet.stream_index == audioStreamIdx) {
            if (avcodec_send_packet(codecContext, &packet) != 0) {
                break;
            }

            memcpy(buffer, pFrame->data[0], dataSize);
        }
        av_packet_unref(&packet);
    }

    // 清理资源
    avcodec_free_context(&codecContext);
    avformat_close_input(&formatCtx);
    av_frame_free(&pFrame);
    av_free(buffer);
}

void PCMDecoder::release() {

}

PCMDecoder pcmDecoder;

extern "C"
JNIEXPORT void JNICALL
Java_com_zxhhyj_ffmpeg_decoder_PCMDecoder_init(JNIEnv *env, jobject, jstring file_path) {
    const char *path = jstringToChar(env, file_path);
    pcmDecoder.init(path);
    delete[] path;
}

extern "C"
JNIEXPORT jbyteArray JNICALL
Java_com_zxhhyj_ffmpeg_decoder_PCMDecoder_decoder(JNIEnv *env, jobject) {
    std::vector<char> result = pcmDecoder.decode();
    jsize resultLen = result.size();
    jbyteArray jArray = env->NewByteArray(resultLen);
    std::vector<jbyte> jbyteResult(result.begin(), result.end());
    env->SetByteArrayRegion(jArray, 0, resultLen, jbyteResult.data());
    return jArray;
}

extern "C"
JNIEXPORT void JNICALL
Java_com_zxhhyj_ffmpeg_decoder_PCMDecoder_release(JNIEnv *, jobject) {
    pcmDecoder.release();
}
extern "C"
JNIEXPORT void JNICALL
Java_com_zxhhyj_ffmpeg_decoder_PCMDecoder_decode(JNIEnv *env, jobject thiz, jobject listener) {

}